#include <objbase.h> 
#include <stdio.h>
#include <atlbase.h>
#include <relogger.h>
#include <windows.h>
#include "CommandLineFilteredTraceEventCallback.h"

HRESULT RelogTraceFile(__in PWSTR inputSessionName, __in PWSTR outputFileName, __in PWSTR imageName)
{
    HRESULT hr;
    CComPtr<ITraceRelogger> relogger;
    CComPtr<ITraceEventCallback> callback;
    TRACEHANDLE traceHandle;

    hr = CoCreateInstance(CLSID_TraceRelogger, NULL, CLSCTX_INPROC_SERVER, __uuidof(ITraceRelogger), (LPVOID*)&relogger);

    if (FAILED(hr)) {
        wprintf(L"Failed to instantiate ITraceRelogger object: 0x%08x.\n", hr);
        return hr;
    }

    hr = relogger->AddLogfileTraceStream(inputSessionName, NULL, &traceHandle);
    if (FAILED(hr)) {
        wprintf(L"Failed to set log file input stream: 0x%08x.\n", hr);
        return hr;
    }

    relogger->SetOutputFilename(outputFileName);
    if (FAILED(hr)) {
        wprintf(L"Failed to set output log file: 0x%08x.\n", hr);
        return hr;
    }

    callback.Attach(new CommandLineFilteredTraceEventCallback(imageName));
    if (callback == NULL) {
        wprintf(L"Failed to allocate callback: 0x%08x.\n", hr);
        return E_FAIL;
    }

    hr = relogger->RegisterCallback(callback);
    if (FAILED(hr)) {
        wprintf(L"Failed to register callback: 0x%08x.\n", hr);
        return hr;
    }

    hr = relogger->ProcessTrace();
    if (FAILED(hr)) {
        wprintf(L"Failed to process trace: 0x%08x.\n", hr);
    }

    return hr;
}

int __cdecl wmain (const int argc, PWSTR* argv)
{
    HRESULT hr;

    if (argc != 5) {
        wprintf(L"Usage: %s <input session name> <directory to store trace files> <filename> \"<commandLineRegex>\"\n", argv[0]);
        return ERROR_INVALID_PARAMETER;
    }

    hr = CoInitializeEx(NULL, COINIT_MULTITHREADED);
    if (FAILED(hr)) {
        wprintf(L"Failed to initialize COM: 0x%08x\n", hr);
        goto Exit;
    }

    /* set workig dir */
    SetCurrentDirectoryW(argv[2]);

    wprintf(L"Starting relogger from session %s to file: %s\\%s\n", argv[1], argv[2], argv[3]);

    hr = RelogTraceFile(argv[1], argv[3], argv[4]);

    if (FAILED(hr))
    {
        wprintf(L"Exiting because Relogging failed.");
        return -1; // for example an external program killed the session
    }

    // uninitializes the com object
    CoUninitialize();
Exit:
    return HRESULT_CODE(hr);
}